import jittor


class FeedForward(jittor.nn.Module):
    def __init__(self, hidden_dim):
        super(FeedForward, self).__init__()
        self.fc_one = jittor.nn.Linear(hidden_dim, hidden_dim // 2, bias=False)
        self.fc_two = jittor.nn.Linear(hidden_dim // 2, hidden_dim, bias=False)
        self.gre = jittor.nn.GELU()

    def execute(self, feed_x):
        feed_x = self.fc_one(feed_x)
        feed_x = self.gre(feed_x)
        feed_x = self.fc_two(feed_x)
        return feed_x


class Block(jittor.nn.Module):
    def __init__(self, hidden_dim):
        super(Block, self).__init__()
        self.layer_nor_x = jittor.nn.LayerNorm(hidden_dim)
        self.layer_nor_p = jittor.nn.LayerNorm(hidden_dim)
        self.layer_a = jittor.nn.Linear(2 * hidden_dim, hidden_dim, bias=True)
        self.layer_b = jittor.nn.Linear(2 * hidden_dim, hidden_dim, bias=True)
        self.layer_c = jittor.nn.Linear(2 * hidden_dim, hidden_dim, bias=True)
        self.layer_nor_a = jittor.nn.LayerNorm(hidden_dim)
        self.layer_nor_b = jittor.nn.LayerNorm(hidden_dim)
        self.layer_nor_c = jittor.nn.LayerNorm(hidden_dim)
        self.max_pool = jittor.nn.MaxPool2d(3)
        self.mlp = FeedForward(hidden_dim)

    def forward(self, p, x):
        x = self.layer_nor_x(x)
        p = self.layer_nor_p(p)
        a = self.layer_a(jittor.concat([x, p], -1))
        b = self.layer_b(jittor.concat([x, p], -1))
        c = self.layer_c(jittor.concat([x, p], -1))
        a = jittor.cumsum(a, 1)
        b = jittor.cumsum(b, 1)
        c = jittor.cumsum(c, 1)
        a = self.layer_nor_a(a)
        b = self.layer_nor_b(b)
        c = self.layer_nor_c(c)
        abc = jittor.concat([a.view([x.shape[0], -1, 1]), b.view([x.shape[0], -1, 1]), c.view([x.shape[0], -1, 1])], -1)
        abc = jittor.concat([abc.view(list(abc.shape) + [1])] * 3, -1)
        x = self.layer_nor_a(x + self.max_pool(abc).view(x.shape))
        #  不加虽然前期低但是网上张
        x=self.mlp(x)+ x
        x = self.layer_nor_b(x)

        return x


class CvFo(jittor.nn.Module):
    def __init__(self, voc_size, hidden_dim, row_layers):
        super(CvFo, self).__init__()
        self.em = jittor.nn.Embedding(voc_size, hidden_dim)
        self.pos = jittor.nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.block = jittor.nn.ModuleList([Block(hidden_dim) for _ in range(row_layers)])
        self.out_layer = jittor.nn.Linear(hidden_dim, voc_size, bias=False)

    def execute(self, sx):
        sx = self.em(sx)
        psx = self.pos(sx)
        for block in self.block:
            sx = block.forward(psx, sx)
        out = self.out_layer(sx)
        return out


if __name__ == '__main__':
    net = CvFo(512, 412, 8)
    net(jittor.randint(0, 123, [2, 15]))
